package com.yuncheng.mcp;

import com.yuncheng.mcp.tool.DynamicTool;
import com.yuncheng.mcp.tool.DynamicToolCallbackProvider;
import com.yuncheng.mcp.tool.SimpleDynamicTool;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.Bean;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Scanner;

@SpringBootApplication
public class McpClientApplication {
    @Autowired
    private List<ToolCallbackProvider> tools;
    @Autowired
    MathService mathService;
    @Autowired
    private DynamicToolCallbackProvider dynamicToolProvider;

    public static void main(String[] args) {
        SpringApplication.run(McpClientApplication.class, args);
    }

    //这样写 就可以既可以用server的tool 也可以使用client的tool
    //>>> 欢迎使用问答系统！输入'exit'退出程序。
    //
    //>>> QUESTION: 你能做什么？
    //
    //>>> ASSISTANT: 我能帮助你完成多种任务，比如进行加法和乘法运算、根据城市名称获取天气预报等。如果你有具体的任务需要我帮忙，比如计算两个数的和或者积，或者是想了解某个城市的天气情况，请告诉我具体需求，我会尽力提供帮助。
    //
    //>>> QUESTION:
    @Bean
    public CommandLineRunner predefinedQuestions(ChatClient.Builder chatClientBuilder,
                                                 ToolCallbackProvider tools,
                                                 ConfigurableApplicationContext context) {
        return args -> {
            // 注册一些动态工具示例
            registerDynamicTools();

            // 打印工具信息
            System.out.println(">>> 动态工具信息：");
            for (DynamicTool tool : dynamicToolProvider.getAllTools()) {
                System.out.println(">>>   工具名称: " + tool.getName());
                System.out.println(">>>     描述: " + tool.getDescription());
                System.out.println(">>>     参数列表:");
                tool.getParameterTypes().forEach((name, type) -> {
                    System.out.println(">>>       - " + name + " (" + type.getSimpleName() + ")");
                });
                System.out.println();
            }

            // 构建ChatClient并注入所有工具
            var chatClient = chatClientBuilder
                    .defaultTools(tools)
                    .build();

            // 创建Scanner对象用于接收用户输入
            Scanner scanner = new Scanner(System.in);

            System.out.println(">>> 欢迎使用问答系统！输入'exit'退出程序。");

            while (true) {
                // 提示用户输入问题
                System.out.print("\n>>> QUESTION: ");
                String userInput = scanner.nextLine();

                // 如果用户输入"exit"，则退出循环
                if ("exit".equalsIgnoreCase(userInput)) {
                    System.out.println(">>> 已退出问答系统。");
                    break;
                }

                // 使用ChatClient与LLM交互
                try {
                    System.out.println("\n>>> ASSISTANT: " + chatClient.prompt(userInput).tools(tools,mathService,dynamicToolProvider).call().content());
                } catch (Exception e) {
                    System.out.println("\n>>> ERROR: 无法处理您的请求，请稍后再试。");
                    e.printStackTrace();
                }
            }

            // 关闭Spring上下文
            context.close();
            scanner.close();
        };
    }

    /**
     * 注册一些动态工具示例
     */
    private void registerDynamicTools() {
        // 示例1: 幂运算工具
        Map<String, Class<?>> powerParams = new HashMap<>();
        powerParams.put("base", Double.class);
        powerParams.put("exponent", Double.class);

        dynamicToolProvider.registerTool(new SimpleDynamicTool(
                "power",
                "计算base的exponent次方",
                "T(Math).pow(#base, #exponent)",
                powerParams
        ));

//        // 示例2: 字符串连接工具
//        Map<String, Class<?>> concatParams = new HashMap<>();
//        concatParams.put("str1", String.class);
//        concatParams.put("str2", String.class);
//
//        dynamicToolProvider.registerTool(new SimpleDynamicTool(
//                "concat",
//                "连接两个字符串",
//                "#str1 + #str2",
//                concatParams
//        ));
    }
}
